import os
"""
数据集整理工具
1、删除txt文件夹中的不匹配编号的文件（没有和图片文件夹编号一一对应）
2、重新按数字顺序排列数据集
"""

############需要修改###################
# 指定图片和文本文件夹的路径
image_folder = '/home/hw/dataset/cone_total_new1/images'
txt_folder = '/home/hw/dataset/cone_total_new1/labels'
#####################################

# 重命名函数
def rename_files(image_folder, txt_folder):
    # 获取图片文件夹中的所有文件名（除去隐藏文件）
    image_files = [f for f in os.listdir(image_folder) if not f.startswith('.')]
    image_files.sort(key=lambda x: int(os.path.splitext(x)[0]))

    # 获取文本文件夹中的所有文件名（除去隐藏文件）
    txt_files = [f for f in os.listdir(txt_folder) if not f.startswith('c')]
    txt_files.sort(key=lambda x: int(os.path.splitext(x)[0]))

    # 输出原始文件数量
    print("原始图片文件夹中共有 {} 个文件，原始文本文件夹中共有 {} 个文件。".format(len(image_files), len(txt_files)))

    # 检查是否存在不匹配的文件，并删除
    mismatched_files = [txt_file for txt_file in txt_files if os.path.splitext(txt_file)[0] + '.jpg' not in image_files]
    if mismatched_files:
        print("以下文本文件编号在图片文件夹中不存在：")
        for file in mismatched_files:
            print(file)
            os.remove(os.path.join(txt_folder, file))
        print("已删除以上文本文件。")

    # 获取文本文件夹中的所有文件名（除去隐藏文件）
    txt_files = [f for f in os.listdir(txt_folder) if not f.startswith('c')]
    txt_files.sort(key=lambda x: int(os.path.splitext(x)[0]))

    # 重命名图片文件
    for i, file in enumerate(image_files):
        _, ext = os.path.splitext(file)
        new_name = str(i) + ext
        # print("重命名图片文件：{} -> {}".format(file, new_name))
        os.rename(os.path.join(image_folder, file), os.path.join(image_folder, new_name))

    # 重命名文本文件
    for i, file in enumerate(txt_files):
        _, ext = os.path.splitext(file)
        new_name = str(i) + ext
        # print("重命名文本文件：{} -> {}".format(file, new_name))
        os.rename(os.path.join(txt_folder, file), os.path.join(txt_folder, new_name))

    # 输出整理后文件夹中的文件数量
    image_files = [f for f in os.listdir(image_folder) if not f.startswith('.')]
    txt_files = [f for f in os.listdir(txt_folder) if not f.startswith('c')]
    print("整理完成，图片文件夹中共有 {} 个文件，文本文件夹中共有 {} 个文件。".format(len(image_files), len(txt_files)))


# 执行整理函数
rename_files(image_folder, txt_folder)